import jwt
from flask import request, jsonify, g, session
from functools import wraps
from db_handler import db_handler  # 假设有一个 db_handler 用于数据库操作
import datetime
from Logger import logger
secret_key = "ckj.1029"  # 请确保将其保存在安全的地方

# 获取当前时间的毫秒级时间戳
def get_current_timestamp_ms():
    return int(datetime.datetime.utcnow().timestamp() * 1000)

def token_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        token = None
        # 从请求头中获取 Token
        if 'Authorization' in request.headers:
            auth_header = request.headers['Authorization']
            parts = auth_header.split()
            if len(parts) == 2 and parts[0] == 'Bearer':
                token = parts[1]
        if not token:
            return jsonify({'message': '缺少Token!'}), 401

        try:
            # 解码 Token
            data = jwt.decode(token, secret_key, algorithms=['HS256'])
            current_user = data['username']
            user_id = data['userid']
            # 将用户信息存储在全局对象 g 中
            g.user = {
                'username': current_user,
                'userid': user_id,
                'roles': data.get('roles', [])
            }

            # 检查会话是否有效
            session_query = """
                SELECT * FROM UserSessions 
                WHERE user_id = %s AND session_token = %s AND expires_at > %s
            """
            current_time = get_current_timestamp_ms()
            session_record = db_handler.execute_query(session_query, (user_id, token, current_time))

            if not session_record:
                # 更新 session_token 为 'expired'
                update_query = """
                              UPDATE UserSessions 
                              SET session_token = 'expired' 
                              WHERE session_token = %s
                          """
                update_result = db_handler.update(update_query, (token,))
                if update_result:
                    logger.info("Token已过期或会话无效，已将 session_token 设置为 'expired'。Token: %s", token)
                else:
                    logger.error("更新 session_token 失败，Token: %s", token)
                return jsonify({'message': '会话无效或已过期!'}), 401

        except jwt.ExpiredSignatureError:
            # Token过期时，更新 session_token 为 'expired'
            update_query = """
                UPDATE UserSessions 
                SET session_token = 'expired' 
                WHERE session_token = %s
            """
            update_result = db_handler.update(update_query, (token,))
            if update_result:
                logger.info("Token已过期，已将 session_token 设置为 'expired'。Token: %s", token)
            else:
                logger.error("更新过期 Token 失败，Token: %s", token)
            return jsonify({'message': 'Token已过期!'}), 401
        except jwt.InvalidTokenError:
            # 更新 session_token 为 'expired'
            update_query = """
                           UPDATE UserSessions 
                           SET session_token = 'expired' 
                           WHERE session_token = %s
                       """
            update_result = db_handler.update(update_query, (token,))
            if update_result:
                logger.info("无效 Token，已将 session_token 设置为 'expired'。Token: %s", token)
            else:
                logger.error("更新无效 Token 失败，Token: %s", token)
            return jsonify({'message': '无效的Token!'}), 401

        return f(*args, **kwargs)

    return decorated